from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget, QToolBar, QComboBox, QLabel, QSpinBox
from PySide6.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

class HorizontalVisualizationWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.current_data = None
        self._init_ui()

    def _init_ui(self):
        layout = QVBoxLayout(self)
        # 工具栏
        toolbar = QToolBar()
        self._setup_toolbar(toolbar)
        layout.addWidget(toolbar)
        # 选项卡控件
        self.tab_widget = QTabWidget()
        self.tab_widget.setTabPosition(QTabWidget.North)
        self.tab_widget.setMovable(True)
        self.tab_widget.setTabsClosable(True)
        self.tab_widget.tabCloseRequested.connect(self._close_tab)
        # 添加三种可视化模式
        self._add_visualization_tab("单井模式", "single_well")
        self._add_visualization_tab("多井对比", "multi_well")
        self._add_visualization_tab("分层对比", "layer_compare")
        layout.addWidget(self.tab_widget)

    def _setup_toolbar(self, toolbar):
        view_type_label = QLabel("视图类型:")
        toolbar.addWidget(view_type_label)
        self.view_type_combo = QComboBox()
        self.view_type_combo.addItems(["曲线图", "柱状图", "散点图", "热力图"])
        self.view_type_combo.currentTextChanged.connect(self._on_view_type_changed)
        toolbar.addWidget(self.view_type_combo)
        toolbar.addSeparator()
        range_label = QLabel("数据范围:")
        toolbar.addWidget(range_label)
        self.range_start = QSpinBox()
        self.range_start.setRange(0, 1000)
        self.range_start.valueChanged.connect(self._on_range_changed)
        toolbar.addWidget(self.range_start)
        range_to_label = QLabel("至")
        toolbar.addWidget(range_to_label)
        self.range_end = QSpinBox()
        self.range_end.setRange(0, 1000)
        self.range_end.setValue(1000)
        self.range_end.valueChanged.connect(self._on_range_changed)
        toolbar.addWidget(self.range_end)

    def _add_visualization_tab(self, title, mode):
        tab = QWidget()
        layout = QVBoxLayout(tab)
        canvas = self._create_plot_canvas(mode)
        layout.addWidget(canvas)
        index = self.tab_widget.addTab(tab, title)
        self.tab_widget.setCurrentIndex(index)
        if not hasattr(self, 'canvases'):
            self.canvases = {}
        self.canvases[mode] = canvas

    def _create_plot_canvas(self, mode):
        fig, ax = plt.subplots(figsize=(8, 4))
        canvas = FigureCanvas(fig)
        self._update_plot_canvas(canvas, mode)
        return canvas

    def update_data(self, filename, data):
        self.current_data = data
        self._update_all_visualizations()

    def _update_all_visualizations(self):
        if self.current_data is None:
            return
        for mode, canvas in self.canvases.items():
            self._update_plot_canvas(canvas, mode)

    def _update_plot_canvas(self, canvas, mode):
        if not isinstance(canvas, FigureCanvas):
            return
        fig = canvas.figure
        ax = fig.axes[0]
        ax.clear()
        if isinstance(self.current_data, pd.DataFrame):
            view_type = self.view_type_combo.currentText()
            if mode == "single_well":
                self._plot_single_well(ax, view_type)
            elif mode == "multi_well":
                self._plot_multi_well(ax, view_type)
            else:
                self._plot_layer_compare(ax, view_type)
        canvas.draw()

    def _plot_single_well(self, ax, view_type):
        if len(self.current_data.columns) >= 2:
            x = self.current_data.iloc[:, 0]
            y = self.current_data.iloc[:, 1]
            if view_type == "曲线图":
                ax.plot(x, y, 'b-', label='数据曲线')
            elif view_type == "柱状图":
                ax.bar(x, y, label='数据分布')
            elif view_type == "散点图":
                ax.scatter(x, y, label='数据点')
            elif view_type == "热力图":
                ax.imshow([y], aspect='auto', cmap='hot')
        ax.legend()
        ax.set_title("单井数据可视化")

    def _plot_multi_well(self, ax, view_type):
        if len(self.current_data.columns) >= 3:
            x = self.current_data.iloc[:, 0]
            for i in range(1, min(4, len(self.current_data.columns))):
                y = self.current_data.iloc[:, i]
                ax.plot(x, y, label=f'井{i}')
        ax.legend()
        ax.set_title("多井数据对比")

    def _plot_layer_compare(self, ax, view_type):
        if len(self.current_data.columns) >= 2:
            data = self.current_data.iloc[:, 1:].T
            ax.boxplot(data)
            ax.set_xticklabels(self.current_data.iloc[:, 0])
        ax.set_title("分层数据对比")

    def _on_view_type_changed(self, view_type):
        self._update_all_visualizations()

    def _on_range_changed(self):
        self._update_all_visualizations()

    def _close_tab(self, index):
        if self.tab_widget.count() > 1:
            self.tab_widget.removeTab(index)

    def update_data(self, well_name, df: pd.DataFrame):
        self.current_well = well_name
        self.data = df

        # 示例：假设需要"斜深"作为X轴
        if '斜深' in df.columns:
            self.plot_widget.plot(
                x=df['斜深'],
                y=df.drop(columns=['斜深'])  # 绘制所有其他列
            )
        else:
            self.show_error("请包含斜深字段")